from typing import Optional, Literal
from pydantic import BaseModel, Field
from fastapi import UploadFile


class ColorizeRequest(BaseModel):
    """Request model for image colorization"""
    
    style: Literal[
        "anime", "realistic", "cartoon", "sketch", "watercolor"
    ] = Field(default="anime", description="Colorization style")
    
    strength: float = Field(
        default=0.7,
        ge=0.1,
        le=1.0,
        description="Colorization strength (0.1-1.0)"
    )
    
    guidance_scale: float = Field(
        default=7.5,
        ge=1.0,
        le=20.0,
        description="Classifier-free guidance scale"
    )
    
    num_inference_steps: int = Field(
        default=20,
        ge=10,
        le=50,
        description="Number of denoising steps"
    )
    
    seed: Optional[int] = Field(
        default=None,
        description="Random seed for reproducible results"
    )
    
    class Config:
        json_schema_extra = {
            "example": {
                "style": "anime",
                "strength": 0.7,
                "guidance_scale": 7.5,
                "num_inference_steps": 20,
                "seed": 42
            }
        }


class BatchColorizeRequest(BaseModel):
    """Request model for batch colorization"""
    
    style: Literal[
        "anime", "realistic", "cartoon", "sketch", "watercolor"
    ] = Field(default="anime", description="Colorization style")
    
    strength: float = Field(
        default=0.7,
        ge=0.1,
        le=1.0,
        description="Colorization strength (0.1-1.0)"
    )
    
    guidance_scale: float = Field(
        default=7.5,
        ge=1.0,
        le=20.0,
        description="Classifier-free guidance scale"
    )
    
    num_inference_steps: int = Field(
        default=20,
        ge=10,
        le=50,
        description="Number of denoising steps"
    )
    
    webhook_url: Optional[str] = Field(
        default=None,
        description="Webhook URL for batch completion notification"
    )
    
    class Config:
        json_schema_extra = {
            "example": {
                "style": "anime",
                "strength": 0.7,
                "guidance_scale": 7.5,
                "num_inference_steps": 20,
                "webhook_url": "https://example.com/webhook"
            }
        }
